# Importing relevant libraries
from sympy import Symbol, Eq, Ge, Abs, Function, Number, sin, cos
from modulus.pdes import PDES
from modulus.variables import Variables
from modulus.solver import Solver
from modulus.dataset import TrainDomain, InferenceDomain
from modulus.data import Inference
import tensorflow as tf
from modulus.sympy_utils.geometry_2d import Rectangle, Line
from modulus.controller import ModulusController
import numpy as np
import sys
# Importing siddharth's code
from kd_tree import kd_Tree # (X, D, N, n, point, Du, U)

# import math
# def get_angle(theta, magnitude):
#     return math.cos(theta)*magnitude, math.sin(theta)*magnitude

# ---------------------------------------------------------------------------------------------------------------------- #
# The PDES class allows you to write
# the equations symbolically in Sympy. This allows users to quickly write their equations in the most natural way possible.
# The Sympy equations are converted to TensorFlow expressions in the back-end and can also be printed to ensure correct
# implementation.

class Poisson_2D(PDES):
    name = 'Poisson_2D'
    def __init__(self):
        x, y = Symbol('x'), Symbol('y')

        alpha = Symbol('alpha')

        input_variables = {'x': x, 'y': y, 'alpha': alpha}

        obstacle_length = 0.10

        phi = Function('phi')(*input_variables)
        u = Function('u')(*input_variables)
        v = Function('v')(*input_variables)
        
        self.equations = Variables()


        self.equations['residual_u'] = u - phi.diff(x)
        self.equations['residual_v'] = v - phi.diff(y)
        self.equations['residual_u_comp'] = u - 10*cos(alpha)
        self.equations['residual_v_comp'] = v - 10*sin(alpha)
        self.equations['residual_obstacle_above'] = v
        self.equations['residual_obstacle_below'] = v
        self.equations['residual_obstacle_wake1_above'] = v # - 10*sin(alpha)*(x)/(3*obstacle_length)
        self.equations['residual_obstacle_wake2_above'] = v # - 10*sin(alpha)*(x)/(3*obstacle_length)
        self.equations['residual_obstacle_wake3_above'] = v # - 10*sin(alpha)*(x)/(3*obstacle_length)

        self.equations['residual_obstacle_wake1_below'] = v # - 10*sin(alpha)*(x)/(3*obstacle_length)
        self.equations['residual_obstacle_wake2_below'] = v # - 10*sin(alpha)*(x)/(3*obstacle_length)
        self.equations['residual_obstacle_wake3_below'] = v # - 10*sin(alpha)*(x)/(3*obstacle_length)

        self.equations['Poisson_2D'] = (phi.diff(x)).diff(x) + (phi.diff(y)).diff(y) # grad^2(phi)

# ---------------------------------------------------------------------------------------------------------------------- #

obstacle_length = 0.10
height = 6*obstacle_length  
width = 6*obstacle_length


rec = Rectangle((-width / 2, -height / 2), (width / 2, height / 2))
obstacle_above = Line((0, 0), (0, obstacle_length), 1)
obstacle_below = Line((0, 0), (0, obstacle_length), 1)

wake1_above = Line((0, -1*obstacle_length), (0, 0), 1) 
wake2_above = Line((0, -2*obstacle_length), (0, -1*obstacle_length), 1) 
wake3_above = Line((0, -3*obstacle_length), (0, -2*obstacle_length), 1) 

wake1_below = Line((0, -1*obstacle_length), (0, 0), 1) 
wake2_below = Line((0, -2*obstacle_length), (0, -1*obstacle_length), 1) 
wake3_below = Line((0, -3*obstacle_length), (0, -2*obstacle_length), 1) 

obstacle_above.rotate(np.pi / 2)
obstacle_below.rotate(np.pi / 2)

wake1_above.rotate(np.pi / 2)
wake2_above.rotate(np.pi / 2)
wake3_above.rotate(np.pi / 2)

wake1_below.rotate(np.pi / 2)
wake2_below.rotate(np.pi / 2)
wake3_below.rotate(np.pi / 2)



geo = rec

x, y, alpha = Symbol('x'), Symbol('y'), Symbol('alpha')

fixed_param_range = {alpha: lambda batch_size: np.full((batch_size, 1), np.random.uniform(-np.pi*10/180, np.pi*10/180))}

#为了在模中生成边界条件，我们需要在

#几何图形的所需边界/曲面，然后为其指定所需的值。

#可以使用boundary_bc函数对边界进行采样。然而，这将对

#几何图形，在本例中为矩形的所有边。几何体的特定边界可以通过以下方式进行子采样

#使用criteria参数为boundary_bc使用特定条件。例如，要对顶墙进行采样，

#条件设置为y=高度/2

#边界条件的所需值在outvar_sympy中作为字典列出。在模数框架中，

#我们将这些变量定义为该字典的键，这些键被转换为计算图中的适当节点。

#使用batch_size_per_area参数指定每个边界上要采样的点数。

#任何变量的权重都可以指定为lambda_sympy参数的输入。添加关键字“lambda_”

#在outvarsympy中指定的键名上指定totalloss函数中的相应权重。

#例如，如果我们想通过速度的大小来加权损失函数，我们可以将权重指定为

#{“lambda_u”：1，“lambda _v”：1、“lambada_phi”：1}

#要求解的方程：

#流体力学的守恒方程在

#几何图形的内部。

#与采样边界类似，我们将使用interior_bc函数对几何体内部的点进行采样。这个

#要求解的方程被指定为outvarsympy的字典输入。

#参数界限决定了变量x和y的值的采样范围。lambda_sympy

#参数用于确定不同损失的权重。在这个问题中，我们在每个点对每个方程进行加权

#通过使用几何图形的“符号距离字段”（SDF），按其与边界的距离计算。这意味着

#与靠近边界的点相比，远离边界的点的权重更高。我们发现这个

#损失函数的加权类型可以加快收敛速度，因为它避免了边界处的不连续性

# ---------------------------------------------------------------------------------------------------------------------- #

class PotentialTrain(TrainDomain):
    def __init__(self, **config):
        super(PotentialTrain, self).__init__()

#############################################################################################
        # I want to make the inlet velocity to be 10.0 m/s with an incidence angle of 4 degrees at the obstacle.
        # the inverse tan(v/u) gives me the required angle of incidence.
        # Drawing the scenario in comments below:
        #      +---------+
        #     /|/     \|/|
        #    //|// --- //|
        #   ///|/////////|
        #  ////+---------+
        #  //////////////
        #  / ////////////
        #    / //////////
        # where / is u + v such that tan-1(v/u) = x degrees(here i kept x as 4).
        u_x = 10*cos(alpha)
        u_y = 10*sin(alpha)
        # flow_rate = u_x*width + u_y*height

        leftWall = geo.boundary_bc(
            outvar_sympy={"residual_u_comp": 0, "residual_v_comp": 0}, # Mimicing the far field conditions "u":u_x , "v": u_y,
            batch_size_per_area=250*2,
            criteria=Eq(x, -width / 2), # As the left wall lies on x = -width/2, we set the criteria to be x = -width/2
            param_ranges ={**fixed_param_range},
            fixed_var=False
        )
        self.add(leftWall, name="LeftWall")

        # outletBC = geo.boundary_bc(
        #     outvar_sympy={"residual_u_comp": 0, "residual_v_comp": 0}, # Mimicing the far field conditions "u":u_x , "v": u_y, 
        #     batch_size_per_area=500*2,
        #     criteria=Ge(y/height+x/width, 1/2),
        #     param_ranges ={**fixed_param_range},
        #     fixed_var=False
        # )
        # self.add(outletBC, name="Outlet")

        topWall = geo.boundary_bc(
            outvar_sympy={"residual_u_comp": 0, "residual_v_comp": 0}, # Mimicing the far field conditions "u":u_x , "v": u_y,
            batch_size_per_area=250*2,
            criteria=Eq(y, height / 2), # As the top wall lies on y = height/2, we set the criteria to be y = height/2
            param_ranges ={**fixed_param_range},
            fixed_var=False
        )
        self.add(topWall, name="TopWall")

        rightWall = geo.boundary_bc(
            outvar_sympy={"residual_u_comp": 0, "residual_v_comp": 0}, # Mimicing the far field conditions "u":u_x , "v": u_y,
            batch_size_per_area=250*2,
            criteria=Eq(x, width / 2), # As the right wall lies on x = width/2, we set the criteria to be x = width/2
            param_ranges ={**fixed_param_range},
            fixed_var=False
        )
        self.add(rightWall, name="RightWall")

        # Bottom Wall
        bottomWall = geo.boundary_bc(
            outvar_sympy={"residual_u_comp": 0, "residual_v_comp": 0}, # Mimicing the far field conditions "u":u_x , "v": u_y,
            batch_size_per_area=250*2,
            criteria=Eq(y, -height / 2), # As the bottom wall lies on y = -height/2, we set the criteria to be y = -height/2
            param_ranges ={**fixed_param_range},
            fixed_var=False            
        )
        self.add(bottomWall, name="BottomWall")

        # obstacleLine Above
        obstacleLineAbove = obstacle_above.boundary_bc(
            outvar_sympy={"u": u_x, 'residual_obstacle_above': 0}, # Setting up the no slip condition for the obstacle. 
            batch_size_per_area=600*2,
            lambda_sympy={"lambda_u": 100, "lambda_residual_obstacle_above": geo.sdf}, # Weights for the loss function.
            param_ranges ={**fixed_param_range},
            # param_ranges ={**y_range_above, **fixed_param_range},
            fixed_var=False            
        )
        self.add(obstacleLineAbove, name="obstacleLineAbove")

        # obstacleLine Below
        obstacleLineBelow = obstacle_below.boundary_bc(
            outvar_sympy={"u": u_x, 'residual_obstacle_below': 0},
            batch_size_per_area=600*2,
            lambda_sympy={"lambda_u": 100, "lambda_residual_obstacle_below": geo.sdf},
            param_ranges ={**fixed_param_range},
            # param_ranges ={**y_range_below, **fixed_param_range},
            fixed_var=False
        )
        self.add(obstacleLineBelow, name="obstacleLineBelow")


        # wakeLine
        # Here we define u = u and v = 0 at the trailing edge of the obstacle(which is at x=0, and v = v at x = right wall). As a linear function for simplicity.
        # As the trailing edge is positioned at {0, 0}, we see the 
        l = lambda x : (x)/(3*obstacle_length) # x = 0 at the trailing edge of the obstacle
        wakeLine1_Above = wake1_above.boundary_bc(
            outvar_sympy={"u": u_x, "v": u_y*l(x), 'residual_obstacle_wake1_above': 0},
            batch_size_per_area=150*2,
            lambda_sympy={"lambda_u": 100, "lambda_v": 100, "lambda_residual_obstacle_wake1_above": geo.sdf},
            param_ranges ={**fixed_param_range},
            # param_ranges ={**y_range_above, **fixed_param_range},
            fixed_var=False            
        )
        self.add(wakeLine1_Above, name="wakeLine1_Above")

        wakeLine2_Above = wake2_above.boundary_bc(
            outvar_sympy={"u": u_x, "v": u_y*l(x), 'residual_obstacle_wake2_above': 0},
            batch_size_per_area=150*2,
            lambda_sympy={"lambda_u": 100, "lambda_v": 100, "lambda_residual_obstacle_wake2_above": geo.sdf},
            param_ranges ={**fixed_param_range},
            # param_ranges ={**y_range_above, **fixed_param_range},
            fixed_var=False
        )
        self.add(wakeLine2_Above, name="wakeLine2_Above")

        wakeLine3_Above = wake3_above.boundary_bc(
            outvar_sympy={"u": u_x, "v": u_y*l(x), 'residual_obstacle_wake3_above': 0},
            batch_size_per_area=150*2,
            lambda_sympy={"lambda_u": 100, "lambda_v": 100, "lambda_residual_obstacle_wake3_above": geo.sdf},
            param_ranges ={**fixed_param_range},
            # param_ranges ={**y_range_above, **fixed_param_range},
            fixed_var=False
        )

        self.add(wakeLine3_Above, name="wakeLine3_Above")

        wakeLine1_Below = wake1_below.boundary_bc(
            outvar_sympy={"u": u_x, "v": u_y*l(x), 'residual_obstacle_wake1_below': 0},
            batch_size_per_area=150*2,
            lambda_sympy={"lambda_u": 100, "lambda_v": 100, "lambda_residual_obstacle_wake1_below": geo.sdf},
            param_ranges ={**fixed_param_range},
            # param_ranges ={**y_range_below, **fixed_param_range},
            fixed_var=False
        )
        self.add(wakeLine1_Below, name="wakeLine1_Below")

        wakeLine2_Below = wake2_below.boundary_bc(
            outvar_sympy={"u": u_x, "v": u_y*l(x), 'residual_obstacle_wake2_below': 0},
            batch_size_per_area=150*2,
            lambda_sympy={"lambda_u": 100, "lambda_v": 100, "lambda_residual_obstacle_wake2_below": geo.sdf},
            param_ranges ={**fixed_param_range},
            # param_ranges ={**y_range_below, **fixed_param_range},
            fixed_var=False
        )
        self.add(wakeLine2_Below, name="wakeLine2_Below")

        wakeLine3_Below = wake3_below.boundary_bc(
            outvar_sympy={"u": u_x, "v": u_y*l(x), 'residual_obstacle_wake3_below': 0},
            batch_size_per_area=150*2,
            lambda_sympy={"lambda_u": 100, "lambda_v": 100, "lambda_residual_obstacle_wake3_below": geo.sdf},
            param_ranges ={**fixed_param_range},
            # param_ranges ={**y_range_below, **fixed_param_range},
            fixed_var=False
        )
        self.add(wakeLine3_Below, name="wakeLine3_Below")

        # interior
        interior = geo.interior_bc(
            outvar_sympy={"Poisson_2D": 0, "residual_u": 0, "residual_v": 0},
            bounds={x: (-width / 2, width / 2), y: (-height / 2, height / 2)}, 
            lambda_sympy={
                "lambda_Poisson_2D": geo.sdf,
                "lambda_residual_u": geo.sdf,
                "lambda_residual_v": geo.sdf,
            },
            batch_size_per_area=4000*2,
            param_ranges ={**fixed_param_range},
            fixed_var=False            
        )
        self.add(interior, name="interior")

        # neighborhood = geo.interior_bc(
        #     outvar_sympy={"Poisson_2D": 0, "residual_u": 0, "residual_v": 0},
        #     bounds={x: (-height / 3, height / 3), y: (-height / 8, height / 8)},
        #     lambda_sympy={
        #         "lambda_Poisson_2D": geo.sdf,
        #         "lambda_residual_u": geo.sdf,
        #         "lambda_residual_v": geo.sdf,
        #     },
        #     batch_size_per_area=2000*2,
        #     param_ranges ={**fixed_param_range},
        #     fixed_var=False            
        # )
        # self.add(neighborhood, name="neighborhood")

# ---------------------------------------------------------------------------------------------------------------------- #

# We will now see how to create an Inference domain to plot the desired variables in the interior
# at a desired point density.
# The LDCInference class can be created by inheriting from the InferenceDomain parent class. The points are again
# sampled in a similar way as done during the definition of LDCTrain and LDCMonitor domains.

class PotentialInference(InferenceDomain):
    def __init__(self, **config):
        super(PotentialInference, self).__init__()
        x, y, alpha = Symbol('x'), Symbol('y'), Symbol('alpha')
        interior = Inference(geo.sample_interior(10000, bounds={x: (-width / 2, width / 2), y: (-height / 2, height / 2)}, param_ranges={alpha: np.pi*(10/180)}), ['u', 'v', 'phi'])
        self.add(interior, name="Inference")

# ---------------------------------------------------------------------------------------------------------------------- #

class PotentialSolver(Solver):
    train_domain = PotentialTrain
    inference_domain = PotentialInference

    def __init__(self, **config):
        super(PotentialSolver, self).__init__(**config)
        self.equations = (
            Poisson_2D().make_node()
        )
        flow_net = self.arch.make_node(
            name="flow_net", inputs=["x", "y", "alpha"], outputs=["u", "v", "phi"]
        )
        self.nets = [flow_net]

#以下函数允许您使用

#插值函数作为矢量点的加权平均值，使用加权数组

#它是相邻点到插值点距离的平方反比的函数。

#片尾：Siddarth Agarwal
    def phi_interpolation(phi, n, weigth_arr, dist):
        interpolated_phi_x = 0 
        phi_x_numer = 0 
        phi_x_denom = 0
        flag_val = 0 
        for i in range(0,n):
            if(dist[i]==0):
                interpolated_phi_x = phi[i] 
                flag_val = 1               
                break
            else:
                phi_x_numer = phi_x_numer + (weigth_arr[i]*phi[i]) 
                phi_x_denom = phi_x_denom + weigth_arr[i]         
        if(flag_val==1):                                           
            return(interpolated_phi_x)
        elif(flag_val==0 and phi_x_denom!=0): 
            return(phi_x_numer/phi_x_denom)

#以下函数使用给定的点云生成子点云

#以及给定点坐标周围的给定边界框。
    def get_sub_pc(self, point, band, x_range, y_range):
        x_range = [[point[0] - x_range*obstacle_length, -obstacle_length][point[0] - x_range*obstacle_length < -obstacle_length], \
            [point[0] + x_range*obstacle_length, width/2][point[0] + x_range*obstacle_length > width/2]] 

        y_range = sorted([0, point[1] + y_range*obstacle_length]) 
        
        sub_pc = [] 
        for i in range(len(band)):
            if x_range[0] <= band[i][0] <= x_range[1] and \
                y_range[0] <= band[i][1] <= y_range[1]:
                sub_pc.append(band[i]) 
        return sub_pc

    def custom_loss(self, domain_invar, pred_domain_outvar, true_domain_outvar, step):
        x_interior = list(tf.make_ndarray(domain_invar['interior']['x'])) + \
            list(tf.make_ndarray(domain_invar['RightWall']['x'])) 
        y_interior = list(tf.make_ndarray(domain_invar['interior']['y'])) + \
            list(tf.make_ndarray(domain_invar['RightWall']['y'])) 

        x_wkeobs_above = domain_invar['obstacleLineAbove']['x'] + \
            domain_invar['wakeLine1_Above']['x'] + \
                domain_invar['wakeLine2_Above']['x'] + \
                    domain_invar['wakeLine3_Above']['x'] 
        y_wkeobs_above = domain_invar['obstacleLineAbove']['y'] + \
            domain_invar['wakeLine1_Above']['y'] + \
                domain_invar['wakeLine2_Above']['y'] + \
                    domain_invar['wakeLine3_Above']['y'] 

        x_wkeobs_below = domain_invar['obstacleLineBelow']['x'] + \
            domain_invar['wakeLine1_Below']['x'] + \
                domain_invar['wakeLine2_Below']['x'] + \
                    domain_invar['wakeLine3_Below']['x'] 
        y_wkeobs_below = domain_invar['obstacleLineBelow']['y'] + \
            domain_invar['wakeLine1_Below']['y'] + \
                domain_invar['wakeLine2_Below']['y'] + \
                    domain_invar['wakeLine3_Below']['y'] 

        # wkeobs_above = np.asarray([x_wkeobs_above, y_wkeobs_above]).T
        # wkeobs_below = np.asarray([x_wkeobs_below, y_wkeobs_below]).T
        
        wkeobs_above = list(zip(x_wkeobs_above, y_wkeobs_above))
        wkeobs_above = [list(t) for t in wkeobs_above] 
        wkeobs_below = list(zip(x_wkeobs_below, y_wkeobs_below))
        wkeobs_below = [list(t) for t in wkeobs_below] 
        interior = list(zip(x_interior, y_interior))
        interior = [list(t) for t in interior] 

        
        band = [] # This is the band of points around the obstacle and wake lines.
        band_range_x = [-obstacle_length, width/2] # This is the range of x values of the band.
        band_range_y = [-0.015, 0.015] # This is the range of y values of the band.

        # The code below filters the interior points to only select those that lie within the range of the band.
        for i in range(len(x_interior)):
            if band_range_x[0] <= x_interior[i] <= band_range_x[1] and \
                band_range_y[0] <= y_interior[i] <= band_range_y[1]:
                band.append([x_interior[i], y_interior[i]])

        # We now have all the points within the band. We need to divide the band into above and below y = 0. 
        # We do this by creating two lists, one for above y = 0 and one for below y = 0.

        band_above = [] # This is the band of points above y = 0.
        band_below = [] # This is the band of points below y = 0.

        for i in range(len(band)):
            if band[i][1] > 0:
                band_above.append(band[i]) # If the y value of the point is above y = 0, we add it to the band above.
            elif band[i][1] < 0:
                band_below.append(band[i]) # If the y value of the point is below y = 0, we add it to the band below.
        
        # We dont concern ourselves with the points that are exactly on the y = 0 line as we are going to add obstacle and wake points 
        # separately in the code below.

        band_above = band_above + wkeobs_above # We add the obstacle and wake points to the band above.
        band_below = band_below + wkeobs_below # We add the obstacle and wake points to the band below.
        
        # Now that we have divided the band into above and below y = 0, we can now start our process of dividing 
        # the point cloud into smaller sub point clouds using my good friend sid's subroutine.

        band_total = band_above + band_below # We combine the band above and below into one list.
        bands = [band_above, band_below] # We create a list of the above and below band to iterate through them later.
        dx = 0.015*obstacle_length # This is the distance between the main point and the constructed points across the x axis.
        dy = 0.015*obstacle_length # This is the distance between the main point and the constructed point across the y axis.
        weights = [] # This is the list of weights of the neighbors used to interpolate the phi value of the constructed points.
        neighbors = [] # These are the neighbors of the constructed points n.
        for i in range(2):
            for j in range(len(bands[i])):
                xfy = [bands[i][j][0] + dx, bands[i][j][1]] # This is the x in front of the original point.
                xby = [bands[i][j][0] - dx, bands[i][j][1]] # This is the x behind the original point.
                xyf = [bands[i][j][0], bands[i][j][1] + (-1)**i*dy] # This is the y next of the original point.
                xy = [bands[i][j][0], bands[i][j][1]] # This is the original point.

                Nxfy = get_sub_pc(xfy, bands[i], 0.3, (-1)**i*0.6) # This is the sub point cloud of the x in front of the original point.
                Nxby = get_sub_pc(xby, bands[i], 0.3, (-1)**i*0.6) # This is the sub point cloud of the x behind the original point.
                Nxy = get_sub_pc(xy, bands[i], 0.3, (-1)**i*0.6) # This is the sub point cloud of the original point.
                Nxyf = Nxy # This is the sub point cloud of the y next of the original point, and is the same as Nxy.

                # We now use sid's subroutine to get the neighbors and weights of the constructed points. 
                # Sid's subroutine uses KD tree to find the neighbors and their corresponding weights for the constructed points.
                # The weights are calculated using the inverse-square of the distance between the constructed point and the neighbors.
                # The neighbors are for now taken as 7 points.
                # The neighbors may be points that are within the band, or outside of it.

                Wxfy, neigh_xfy = kd_Tree(Nxfy, 2, len(Nxfy), 7, xfy)
                Wxby, neigh_xby = kd_Tree(Nxby, 2, len(Nxby), 7, xby)
                Wxyf, neigh_xyf = kd_Tree(Nxyf, 2, len(Nxyf), 7, xyf)
                # Wxy, neigh_xy = kd_Tree(Nxy, 2, len(Nxy), 7, xy)

                weights.append([Wxfy, Wxby, Wxyf]) # We add the weights of the neighbors to the weights list.
                neighbors.append([neigh_xfy, neigh_xby, neigh_xyf, xy]) # We add the neighbors to the neighbors list.
        
        # In the code above, we need 4 points: two points on the x-axis for central differentiation and one 
        # point on the y-axis for backward differentiation. We then use these points to find the sub point cloud around them. 
        # Using these sub point clouds we can then calculate the weights and neighbors for each point. 
        # The weights and neighbors are calculated using the kdTree function(sid's subroutine). 
        # We then store them in the weights and neighbors list, which will be used later. 
        # For now each entry in weight and neighbor corresponds to information about 4 points: xfy, xby, xyf, xy. 
        # Since we specified the neighbors as 7, we will have 7 dimensional vector for each of theses points, 
        # making it a total of 4 times 7 = 28 entries per point (x,y).

        u_band = [] # This is the list that stores the u values within the band initially. 
        v_band = [] # This is the list that stores the v values within the band initially.
        for i in range(len(neighbors[0])):
            # xfy
            x_xfy = [x for x,y in neighbors[i][0]] # We create a list of the x values of the neighbors of xfy.
            y_xfy = [y for x,y in neighbors[i][0]] # We create a list of the y values of the neighbors of xfy.

            # xby
            x_xby = [x for x,y in neighbors[i][1]] # We create a list of the x values of the neighbors of xby.
            y_xby = [y for x,y in neighbors[i][1]] # We create a list of the y values of the neighbors of xby.

            # xyf
            x_xyf = [x for x,y in neighbors[i][2]] # We create a list of the x values of the neighbors of xyf.
            y_xyf = [y for x,y in neighbors[i][2]] # We create a list of the y values of the neighbors of xyf.

            # phi xfy
            phi_xfy = self.nets[0].evaluate({'x': x_xfy, 'y': y_xfy})['phi'] # We evaluate the phi values of the neighbors of xfy using the neural network.

            # phi xby
            phi_xby = self.nets[0].evaluate({'x': x_xby, 'y': y_xby})['phi'] # We evaluate the phi values of the neighbors of xby using the neural network.

            # phi xyf
            phi_xyf = self.nets[0].evaluate({'x': x_xyf, 'y': y_xyf})['phi'] # We evaluate the phi values of the neighbors of xyf using the neural network.

            # phi xy
            phi_xy = self.nets[0].evaluate({'x': xy[0], 'y': xy[1]})['phi'] # We evaluate the phi values of the original point using the neural network.

            # Interpolating phi for xfy, xby, xyf using phi_interpolation()

            phi_xfy = phi_interpolation(phi_xfy, len(phi_xfy), weights[i][0]) # We interpolate the phi values of the neighbors of xfy using the weighted average.

            phi_xby = phi_interpolation(phi_xby, len(phi_xby), weights[i][1]) # We interpolate the phi values of the neighbors of xby using the weighted average.

            phi_xyf = phi_interpolation(phi_xyf, len(phi_xyf), weights[i][2]) # We interpolate the phi values of the neighbors of xyf using the weighted average.

            # Calculating u and v
            u_band.append([xy[0], xy[1]], [(phi_xfy - phi_xby)/(2*dx)]) # We calculate the u value of the original point using the central difference formula.
            v_band.append([xy[0], xy[1]], [abs(phi_xyf - phi_xy)/(dy)]) # We calculate the v value of the original point using the forward/backward difference formula.

        # We now have the u and v values for each point in the point cloud. We now need to find the u and v values outside the band.
        # We do this by finding the points that are outside the band and then using the neural network to evaluate them.

        # We first find the points that are outside the band.

        x_outside = [x for x,y in interior if (x,y) not in band_total]
        y_outside = [y for x,y in interior if (x,y) not in band_total]
                
        # We now use the neural network(self.nets[0]) to evaluate the points that are outside the band.
        phi_outside = self.nets[0].evaluate({'x': x_outside, 'y': y_outside})['phi']

        u_outside = tf.gradients(phi_outside, x_outside)[0] # We calculate the u values of the points outside the band using d(phi)/dx (automatic differentiation).
        v_outside = tf.gradients(phi_outside, y_outside)[0] # We calculate the v values of the points outside the band using d(phi)/dy (automatic differentiation).

    @classmethod
    def update_defaults(cls, defaults):
        defaults.update(
            {
                "network_dir": "./network_checkpoint_potential_flow_2d",
                "decay_steps": 4000,
                "max_steps": 400000,
                "layer_size": 100,
            }
        )

# ---------------------------------------------------------------------------------------------------------------------- #

if __name__ == "__main__":
    ctr = ModulusController(PotentialSolver)
    ctr.run()